#include <jni.h>
#include <opus/opus.h>
#include <stdlib.h>
#include <string.h>
#include "opus_wrapper.h"

#define FRAME_SIZE 320  // 16kHz * 20ms = 320 samples
#define MAX_PACKET_SIZE 1024

typedef struct {
    OpusEncoder* encoder;
    OpusDecoder* decoder;
} OpusCodecContext;

JNIEXPORT jlong JNICALL Java_com_example_OpusCodec_initEncoder
  (JNIEnv *env, jobject obj, jint sample_rate, jint channels, jint bitrate) {
    int err;
    OpusEncoder* encoder = opus_encoder_create(sample_rate, channels, OPUS_APPLICATION_RESTRICTED_LOWDELAY, &err);
    if (err != OPUS_OK) return 0;
    
    // 配置编码参数
    opus_encoder_ctl(encoder, OPUS_SET_BITRATE(bitrate));
    opus_encoder_ctl(encoder, OPUS_SET_VBR(0));         // 禁用VBR，启用CBR
    opus_encoder_ctl(encoder, OPUS_SET_COMPLEXITY(0)); 
    
    OpusCodecContext* ctx = malloc(sizeof(OpusCodecContext));
    ctx->encoder = encoder;
    ctx->decoder = NULL;
    return (jlong)ctx;
}

JNIEXPORT jlong JNICALL Java_com_example_OpusCodec_initDecoder
  (JNIEnv *env, jobject obj, jint sample_rate, jint channels) {
    int err;
    OpusDecoder* decoder = opus_decoder_create(sample_rate, channels, &err);
    if (err != OPUS_OK) return 0;
    
    OpusCodecContext* ctx = malloc(sizeof(OpusCodecContext));
    ctx->decoder = decoder;
    ctx->encoder = NULL;
    return (jlong)ctx;
}

JNIEXPORT jbyteArray JNICALL Java_com_example_OpusCodec_encodeFrame
  (JNIEnv *env, jobject obj, jlong ctx_ptr, jbyteArray pcm_data) {
    OpusCodecContext* ctx = (OpusCodecContext*)ctx_ptr;
    jbyte* pcm_bytes = (*env)->GetByteArrayElements(env, pcm_data, NULL);
    jsize len = (*env)->GetArrayLength(env, pcm_data);
    
    // 确保输入数据长度正确（每个short占2字节）
    if (len != FRAME_SIZE * sizeof(short)) {
        (*env)->ReleaseByteArrayElements(env, pcm_data, pcm_bytes, JNI_ABORT);
        return NULL;
    }

    unsigned char out[MAX_PACKET_SIZE];
    // 将byte数组直接转换为short指针（注意字节顺序应与Java端一致）
    short* pcm = (short*)pcm_bytes;
    int ret = opus_encode(ctx->encoder, pcm, FRAME_SIZE, out, MAX_PACKET_SIZE);
    
    (*env)->ReleaseByteArrayElements(env, pcm_data, pcm_bytes, JNI_ABORT);
    
    if (ret < 0) return NULL;
    
    jbyteArray result = (*env)->NewByteArray(env, ret);
    (*env)->SetByteArrayRegion(env, result, 0, ret, (jbyte*)out);
    return result;
}

JNIEXPORT jbyteArray JNICALL Java_com_example_OpusCodec_decodeFrame
  (JNIEnv *env, jobject obj, jlong ctx_ptr, jbyteArray opus_data) {
    OpusCodecContext* ctx = (OpusCodecContext*)ctx_ptr;
    jbyte* data = (*env)->GetByteArrayElements(env, opus_data, NULL);
    jsize len = (*env)->GetArrayLength(env, opus_data);
    
    short pcm[FRAME_SIZE];
    int ret = opus_decode(ctx->decoder, (unsigned char*)data, len, pcm, FRAME_SIZE, 0);
    
    (*env)->ReleaseByteArrayElements(env, opus_data, data, JNI_ABORT);
    
    if (ret < 0) return NULL;
    
    // 将short数组转换为byte数组（注意字节顺序）
    jbyteArray result = (*env)->NewByteArray(env, FRAME_SIZE * sizeof(short));
    jbyte* bytes = (*env)->GetByteArrayElements(env, result, NULL);
    memcpy(bytes, pcm, FRAME_SIZE * sizeof(short));
    (*env)->ReleaseByteArrayElements(env, result, bytes, 0);
    
    return result;
}

JNIEXPORT void JNICALL Java_com_example_OpusCodec_destroy
  (JNIEnv *env, jobject obj, jlong ctx_ptr) {
    OpusCodecContext* ctx = (OpusCodecContext*)ctx_ptr;
    if (ctx->encoder) opus_encoder_destroy(ctx->encoder);
    if (ctx->decoder) opus_decoder_destroy(ctx->decoder);
    free(ctx);
}